package my;

import java.nio.ByteBuffer;
import java.util.Queue;
import java.util.concurrent.ArrayBlockingQueue;

/**
 * @author fengbo
 */
public class ByteBufferAllocator {

    private final Queue<ByteBuffer>[] caches;

    public ByteBufferAllocator() {
        caches = new ArrayBlockingQueue[10];
        for (int i = 0; i < caches.length; i++) {
            caches[i] = new ArrayBlockingQueue<>(20);
        }
    }

    public ByteBuffer allocate(int limit) {
        int size = limit / 1024;
        if (size >= 10) {
            return null;
        }
        Queue<ByteBuffer> queue = caches[size];
        ByteBuffer result = queue.poll();
        if (result == null) {
            result = ByteBuffer.allocateDirect(limit);
        }
        return result;
    }

    public void release(ByteBuffer byteBuffer, int limit) {
        int size = limit / 1024;
        if (size < 10) {
            Queue<ByteBuffer> queue = caches[size];
            byteBuffer.clear();
            queue.offer(byteBuffer);
        }
    }
}
